from models import Stock, Industry
import akshare as ak
import pandas as pd
import sqlite3 

def save_to_sqlite():
    try:
        df = ak.stock_info_a_code_name()
        print(df)

        conn = sqlite3.connect('db/stock.db')
        df.to_sql('stocks', conn, if_exists='replace', index=False)
        conn.close()
        return True
    except Exception as e:
        print(f"保存数据失败: {e}")
        return False

def load_data_from_akshare():
    """从Akshare API加载股票数据"""
    try:
        # 获取A股上市公司基本信息
        stock_info_a_df = ak.stock_info_a_code_name()
        
        # 获取行业分类数据
        industry_data = ak.stock_sector_spot(indicator="板块名称")
        
        # 创建行业字典
        industries = {}
        stocks = []
        
        # 处理每个行业
        for _, row in industry_data.iterrows():
            industry_name = row['板块名称']
            if '行业' not in industry_name:
                continue  # 只处理行业板块
                
            # 清理行业名称
            industry_name = industry_name.replace('行业', '')
            
            # 获取该行业的股票
            industry_stocks = ak.stock_board_industry_cons(symbol=row['板块名称'])
            
            # 创建行业对象
            if industry_name not in industries:
                industries[industry_name] = Industry(industry_name)
            
            # 处理该行业的每只股票
            for _, stock_row in industry_stocks.iterrows():
                stock_code = stock_row['代码']
                stock_name = stock_row['名称']
                
                # 获取股票最新价格
                try:
                    stock_price_df = ak.stock_zh_a_spot()
                    price_data = stock_price_df[stock_price_df['代码'] == stock_code]
                    price = float(price_data['最新价'].values[0]) if not price_data.empty else 0.0
                except:
                    price = 0.0  # 如果获取价格失败，设为0
                
                # 创建股票对象
                stock = Stock(stock_code, stock_name, price, industry_name)
                stocks.append(stock)
                industries[industry_name].add_stock(stock)
        
        return stocks, industries
    except Exception as e:
        print(f"从Akshare加载数据时出错: {e}")

    """加载示例股票数据"""
    # 创建示例股票数据
    sample_stocks = [
        Stock("000001", "平安银行", 10.5, "银行业"),
        Stock("000002", "万科A", 15.2, "房地产"),
        Stock("000063", "中兴通讯", 30.8, "通信设备"),
        Stock("000069", "华侨城A", 5.7, "旅游业"),
        Stock("000100", "TCL科技", 4.2, "电子设备"),
        Stock("000157", "中联重科", 6.8, "机械设备"),
        Stock("000333", "美的集团", 45.6, "家用电器"),
        Stock("000338", "潍柴动力", 12.3, "汽车制造"),
        Stock("000568", "泸州老窖", 150.2, "白酒"),
        Stock("000596", "古井贡酒", 135.7, "白酒"),
        Stock("000651", "格力电器", 38.5, "家用电器"),
        Stock("000661", "长春高新", 180.3, "医药制造"),
        Stock("000725", "京东方A", 3.8, "电子设备"),
        Stock("000858", "五粮液", 160.5, "白酒"),
        Stock("000876", "新希望", 8.9, "农牧业"),
        Stock("001979", "招商蛇口", 14.2, "房地产"),
        Stock("002142", "宁波银行", 25.6, "银行业"),
        Stock("002304", "洋河股份", 120.8, "白酒"),
        Stock("002415", "海康威视", 32.1, "电子设备"),
        Stock("002714", "牧原股份", 45.3, "农牧业"),
    ]
    
    # 按行业分类
    industries = {}
    for stock in sample_stocks:
        if stock.industry_name not in industries:
            industries[stock.industry_name] = Industry(stock.industry_name)
        industries[stock.industry_name].add_stock(stock)
    
    # 返回所有股票和行业分类
    return sample_stocks, industries

def load_data_from_csv(file_path):
    """从CSV文件加载股票数据"""
    # 这里可以实现从CSV文件加载数据的功能
    # 暂时返回示例数据
    return load_sample_data()


def get_stock_total():
    """获取A股总数"""
    try:
        conn = sqlite3.connect('db/stock.db')
        df = pd.read_sql('SELECT * FROM stocks', conn)
        conn.close()

        return len(df)
    except Exception as e:
        print(f"Akshare获取总数失败: {e}")
        return 0

def get_stock_list():
    """获取A股列表"""
    try:
        # 从数据库获取股票列表
        conn = sqlite3.connect('db/stock.db')
        df = pd.read_sql('SELECT * FROM stocks', conn)
        conn.close()
        return df.to_dict(orient='records')

    except Exception as e:
        print(f"Akshare获取列表失败: {e}")
        return []